from fastapi import Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp
import json
from typing import Any
import logging

logger = logging.getLogger(__name__)


class ResponseFormatterMiddleware(BaseHTTPMiddleware):
    """统一响应格式化中间件"""

    def __init__(self, app: ASGIApp):
        super().__init__(app)

    async def dispatch(self, request: Request, call_next):
        # 排除的路径（不需要格式化的接口）
        excluded_paths = [
            "/docs",
            "/redoc",
            "/openapi.json",
            "/uploads",
            "/health",
            "/favicon.ico",
        ]

        # 检查是否是排除的路径
        for excluded_path in excluded_paths:
            if request.url.path.startswith(excluded_path):
                return await call_next(request)

        # 处理请求
        response = await call_next(request)

        # 只处理成功的JSON响应
        if response.status_code == 200 and "application/json" in response.headers.get(
            "content-type", ""
        ):

            # 读取响应体
            body = b""
            async for chunk in response.body_iterator:
                body += chunk

            try:
                # 解析JSON数据
                data = json.loads(body.decode())

                # 检查是否已经是统一格式
                if isinstance(data, dict) and "code" in data and "message" in data:
                    # 已经是统一格式，直接返回
                    formatted_data = data
                else:
                    # 转换为统一格式
                    formatted_data = {"code": 200, "message": "success", "data": data}

                # 创建新的响应头，排除Content-Length让FastAPI自动计算
                new_headers = {}
                for key, value in response.headers.items():
                    if key.lower() not in ["content-length", "transfer-encoding"]:
                        new_headers[key] = value

                # 返回格式化后的响应
                return JSONResponse(
                    content=formatted_data,
                    status_code=200,
                    headers=new_headers,
                )

            except Exception as e:
                logger.error(f"Response formatting error: {e}")
                # 如果格式化失败，返回原响应内容
                try:
                    original_data = json.loads(body.decode()) if body else {}
                    new_headers = {}
                    for key, value in response.headers.items():
                        if key.lower() not in ["content-length", "transfer-encoding"]:
                            new_headers[key] = value

                    return JSONResponse(
                        content=original_data,
                        status_code=response.status_code,
                        headers=new_headers,
                    )
                except:
                    # 最后的备选方案，返回空对象
                    return JSONResponse(
                        content={
                            "code": 500,
                            "message": "Response formatting failed",
                            "data": None,
                        },
                        status_code=500,
                    )

        return response


def format_success_response(
    data: Any = None, message: str = "success", code: int = 200
) -> dict:
    """格式化成功响应的辅助函数"""
    return {"code": code, "message": message, "data": data}


def format_error_response(
    message: str = "error", code: int = 400, data: Any = None
) -> dict:
    """格式化错误响应的辅助函数"""
    return {"code": code, "message": message, "data": data}
